# 最接近原点的K个点： https://leetcode-cn.com/problems/k-closest-points-to-origin/submissions/

class Solution:
    def kClosest(self, points: list, k: int) -> list:
        """
            这道题和 215号算法题，数组中第k个最大元素
            347 前K个高频元素 一模一样~~~~
        """
        def compare(n1, n2):
            """ 定义比较两个点 大小的函数 """
            return n1[0] ** 2 + n1[1] ** 2 < n2[0] ** 2 + n2[1] ** 2

        # 剩下的步骤使用最小堆~ 和上面两题如出一辙！
        def heapqAdjust(nums):
            """ 调整堆， 比较大小需要使用 freq的value """
            n = len(nums)
            i = 0
            end = n - 1
            rc = nums[0]
            while i * 2 + 1 <= end:
                j = i * 2 + 1 
                if j < end and compare(nums[j], nums[j + 1]): j += 1 
                if compare(nums[j], rc): break 
                nums[i] = nums[j]
                i = j
            
            nums[i] = rc
        
        def heapqPush(nums, num):
            """ 添加元素进堆 """
            nums.append(num)
            n = len(nums)
            i = n - 1
            while (i - 1) // 2 >= 0:
                j = (i - 1) // 2
                if not compare(nums[j], num): break
                nums[i] = nums[j]
                i = j

            nums[i] = num

        # 3. 开始逻辑
        heap = []
        for num in points:
            if len(heap) < k:
                heapqPush(heap, num)

            elif compare(num, heap[0]):
                heap[0] = num
                heapqAdjust(heap)

        return heap 



# 官方实现解法
import heapq
class Solution:
    def kClosest(self, points: list, k: int) -> list:
        q = [(-x ** 2 - y ** 2, i) for i, (x, y) in enumerate(points[:k])]
        heapq.heapify(q)
        
        n = len(points)
        for i in range(k, n):
            x, y = points[i]
            dist = -x ** 2 - y ** 2
            heapq.heappushpop(q, (dist, i))
        
        ans = [points[identity] for (_, identity) in q]
        return ans
